package cryptohelp

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
)

type RsaHelp struct {
}

// Encrypt 加密
func (c *RsaHelp) Encrypt(publicKey *rsa.PublicKey, plaintext []byte) ([]byte, error) {
	return rsa.EncryptOAEP(sha256.New(), rand.Reader, publicKey, plaintext, nil)
}

// EncryptStr 返回加密字符串
func (c *RsaHelp) EncryptStr(publicKey *rsa.PublicKey, plaintext []byte) (string, error) {
	bytes, err := c.Encrypt(publicKey, plaintext)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(bytes), nil
}

// Decrypt 解密
func (c *RsaHelp) Decrypt(privateKey *rsa.PrivateKey, ciphertext []byte) ([]byte, error) {
	plaintext, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, privateKey, ciphertext, nil)
	if err != nil {
		return nil, err
	}
	return plaintext, nil
}

// DecryptStr 解密字符串
func (c *RsaHelp) DecryptStr(privateKey *rsa.PrivateKey, str string) ([]byte, error) {
	// 将字符串转换为字节切片
	bytes, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return nil, err
	}
	return c.Decrypt(privateKey, bytes)
}

// ConvertPrivateKeyToPEM 将 RSA 私钥转换为 PEM 格式的字符串
func (c *RsaHelp) ConvertPrivateKeyToPEM(privateKey *rsa.PrivateKey, isPKCS8 bool) (string, error) {
	// 将私钥转换为 DER 格式
	var privateKeyDER []byte
	var err error
	if isPKCS8 {
		privateKeyDER, err = x509.MarshalPKCS8PrivateKey(privateKey)
		if err != nil {
			return "", err
		}
	} else {
		privateKeyDER = x509.MarshalPKCS1PrivateKey(privateKey)
	}

	var typeStr string
	if isPKCS8 {
		typeStr = "PRIVATE KEY"
	} else {
		typeStr = "RSA PRIVATE KEY"
	}
	// 将 DER 格式的私钥编码为 PEM 格式
	privateKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  typeStr,       // PEM 块类型
		Bytes: privateKeyDER, // DER 编码的私钥数据
	})

	return string(privateKeyPEM), nil
}

// ConvertPublicKeyToPEM 将 RSA 公钥钥转换为 PEM 格式的字符串
func (c *RsaHelp) ConvertPublicKeyToPEM(publicKey *rsa.PublicKey, isPKCS8 bool) (string, error) {
	// 将私钥转换为 DER 格式
	var publicKeyDER []byte
	var err error
	if isPKCS8 {
		publicKeyDER, err = x509.MarshalPKIXPublicKey(publicKey)
		if err != nil {
			return "", err
		}
	} else {
		publicKeyDER = x509.MarshalPKCS1PublicKey(publicKey)
	}

	var typeStr string
	if isPKCS8 {
		typeStr = "PUBLIC KEY"
	} else {
		typeStr = "RSA PUBLIC KEY"
	}
	// 将 DER 格式的私钥编码为 PEM 格式
	publicKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  typeStr,      // PEM 块类型
		Bytes: publicKeyDER, // DER 编码的私钥数据
	})

	return string(publicKeyPEM), nil
}

func (c *RsaHelp) LoadPemPrivateKey(pemStr string) (*rsa.PrivateKey, error) {

	block, _ := pem.Decode([]byte(pemStr))
	if block == nil {
		return nil, errors.New("pem decode failed")
	}

	priKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		// 如果不是 PKCS#8 格式，尝试解析 PKCS#1
		privateKeyRSA, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		priKey = privateKeyRSA
	}

	rsaPriKey, ok := priKey.(*rsa.PrivateKey)
	if !ok {
		return nil, errors.New("私钥不是 RSA 类型")
	}

	return rsaPriKey, nil
}

func (c *RsaHelp) LoadPemPublicKey(pemStr string) (*rsa.PublicKey, error) {

	block, _ := pem.Decode([]byte(pemStr))
	if block == nil {
		return nil, errors.New("pem decode failed")
	}

	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		// 如果不是 PKCS#8 格式，尝试解析 PKCS#1
		pubKeyRSA, err := x509.ParsePKCS1PublicKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		pubKey = pubKeyRSA
	}

	rsaPubKey, ok := pubKey.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("公钥不是 RSA 类型")
	}

	return rsaPubKey, nil
}
